from functools import reduce
import math
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
from tensorflow.keras import layers
from losses_utils import *

@tf.function
def yolo_loss_on_size(labels,predns,grid,grid_anchors,xyxy=False,GIoU=False,DIoU=False,CIoU=False,nc=80,na=3,label_smoothing=0.0001,
                      filter=True,weight=False):
    """
        predns为张量, (B,nx*ny,3,5+nc)
        labels为张量, (B,nx*ny,3,5+nc) [xc,yx,w,h,r,c...]
        #! loop 采用for循环计算, 还是张量整体计算
        #! filter 将正样本分离出来计算, 还是不分离计算
        #! weight 按正样本数量计算, 还是按权重计算
    """
    if not isinstance(labels,tf.Tensor): labels = tf.convert_to_tensor(labels,dtype=tf.float32)
    else: labels = tf.cast(labels,dtype=tf.float32)

    #! 还原pbox 从 偏移量 到 具体坐标
    pxy = tf.nn.sigmoid(predns[:,:,:,0:2]) * 2 - 0.5 + tf.cast(grid,dtype=predns.dtype)     # (B,ny*nx,3,2) + (ny*nx,3,2)
    pwh = (tf.nn.sigmoid(predns[:,:,:,2:4]) * 2) ** 2 * tf.cast(grid_anchors,dtype=predns.dtype) # (B,ny*nx,3,2) * (3,2)
    predns = tf.concat([pxy,pwh,predns[...,4:]],axis=-1)

    #! 初始化变量
    lbox = 0.0
    lcls = 0.0
    lobj = 0.0

    #! 进行损失计算
    for adx in range(na):
        indices = tf.equal(labels[...,adx,4],1.0) #! (B,nx*ny)
        if tf.reduce_any(indices): #! 如果存在匹配
            masks = tf.cast(indices,labels.dtype)
            instances = tf.maximum(tf.reduce_sum(masks),1.0)

            t = tf.gather_nd(labels[:,:,adx],tf.where(indices))
            p = tf.gather_nd(predns[:,:,adx],tf.where(indices))
            tf.print("%d t:"%adx)
            tf.print(t[:5],summarize=-1)
            tf.print("%d p:"%adx)
            tf.print(p[:5],summarize=-1)

            t0 = tf.gather_nd(labels[:,:,adx],tf.where(~indices))
            p0 = tf.gather_nd(predns[:,:,adx],tf.where(~indices))
            tf.print("%d t0:"%adx)
            tf.print(t0[:5],summarize=-1)
            tf.print("%d p0:"%adx)
            tf.print(p0[:5],summarize=-1)

            _iou = bbox_iou(labels[:,:,adx,:4],predns[:,:,adx,:4],xyxy=xyxy,GIoU=GIoU,DIoU=DIoU,CIoU=CIoU)
            # _cls = keras.losses.categorical_crossentropy(
            _cls = keras.losses.binary_crossentropy(
                (1 - label_smoothing) * labels[:,:,adx,5:] + label_smoothing / nc,
                predns[:,:,adx,5:],from_logits=True
            )
            obij = tf.maximum(_iou,0) * masks
            prob = tf.reduce_max(predns[:,:,adx,5:],axis=-1) #! (B,nx*ny)

            if filter:
                _box = tf.reduce_mean(1-tf.gather_nd(_iou,tf.where(indices)))
                _cls = tf.reduce_mean(tf.gather_nd(_cls,tf.where(indices)))
            else:
                if weight:
                    weight_mask = labels[:,:,adx,4]
                    _box = tf.reduce_sum(weight_mask * (1-_iou)) / instances
                    _cls = tf.reduce_sum(weight_mask * _cls) / instances
                else:
                    _box = tf.reduce_sum(masks * (1-_iou)) / instances
                    _cls = tf.reduce_sum(masks * _cls) / instances
            
            # _obj = keras.losses.binary_crossentropy(obij,prob * predns[:,:,adx,4],from_logits=True)
            _obj = keras.losses.binary_crossentropy(obij,predns[:,:,adx,4],from_logits=True)
            _obj = tf.reduce_mean(_obj)
            
            lbox += _box
            lcls += _cls
            lobj += _obj

    return lbox, lcls, lobj

@tf.function
def yolo_loss(labels,predns,grids,grid_anchors,xyxy=False,GIoU=False,DIoU=False,CIoU=False,nc=80,na=3,label_smoothing=0.0001,
              filter=True,weight=False):
    """
        predns为张量, (B,nx*ny,3,5+nc)
        labels为张量, (B,nx*ny,3,5+nc) [xc,yx,w,h,r,c...]
        #! loop为False的情况下, filter失效
        #! loop为True, filter为True的情况下, weight失效
    """
    assert len(labels) == len(predns)

    lbox = 0.0
    lcls = 0.0
    lobj = 0.0
    balance = [4.0,1.0,0.4]

    for ldx,(label,predn) in enumerate(zip(labels,predns)):
        tf.print("ldx: ",ldx)
        box_loss, cls_loss, obj_loss = yolo_loss_on_size(
            label,predn,grids[ldx],grid_anchors[ldx],xyxy=xyxy,GIoU=GIoU,DIoU=DIoU,CIoU=CIoU,nc=nc,na=na,label_smoothing=label_smoothing,
            filter=filter,weight=weight)
        lbox += box_loss
        lcls += cls_loss
        lobj += obj_loss * balance[ldx]

    #! 使用model.compile中的loss_fn功能会把损失进行reduce化, 默认reduce函数是reduce_mean
    #! 可以使用reduce_sum, reduce_none, reduce_sum_over_batch进行替换, 但需要使用类处理
    #! 若直接使用loss_fn传入model.compile, 则会直接reduce_auto(默认是reduce_mean)
    
    return lbox, lcls, lobj